# models/stock_kline.py
import pandas as pd
import sqlalchemy as sa
from datetime import datetime
from typing import Optional

from infrastructure.database.database_context import DatabaseContext
from infrastructure.model.base_model import BaseModel


class StockKline(BaseModel):
    __tablename__ = 'stock_kline'

    ts_code = sa.Column(sa.CHAR(9), primary_key=True)
    trade_date = sa.Column(sa.DateTime(), primary_key=True)
    open = sa.Column(sa.Numeric(8, 2))
    high = sa.Column(sa.Numeric(8, 2))
    low = sa.Column(sa.Numeric(8, 2))
    close = sa.Column(sa.Numeric(8, 2))
    volume = sa.Column(sa.BigInteger)

    @classmethod
    def get_stock_data(
            cls,
            ts_code: str,
            start_date: Optional[datetime] = None,
            end_date: Optional[datetime] = None
    ) -> pd.DataFrame:
        """获取股票数据"""
        session = DatabaseContext.get_session()

        query = session.query(cls).filter(cls.ts_code == ts_code)

        if start_date:
            query = query.filter(cls.trade_date >= start_date)
        if end_date:
            query = query.filter(cls.trade_date <= end_date)

        query = query.order_by(cls.trade_date)
        results = query.all()

        return cls.to_dataframe(results)

    @classmethod
    def get_latest_date(cls, ts_code: str) -> Optional[datetime]:
        """Get the latest date for a symbol from database"""
        session = DatabaseContext.get_session()
        query = session.query(cls.trade_date).filter(cls.ts_code == ts_code).order_by(cls.trade_date.desc()).first()
        return query.trade_date if query else None
